import os
import time
import random
import threading
import datetime
import argparse

def generate_test_data(size):
    """生成随机测试数据（模拟视频流）"""
    return os.urandom(size)

def test_recording(usb_version, test_dir, duration_seconds=300, write_rate_mbps=20):
    """
    测试U盘在指定接口下的录制性能
    
    参数:
    - usb_version: 接口版本描述（如"USB3.0"或"USB2.0"）
    - test_dir: 测试目录路径（U盘根目录）
    - duration_seconds: 测试持续时间（秒）
    - write_rate_mbps: 目标写入速率（Mbps）
    """
    # 计算每帧大小和写入间隔（模拟30fps视频）
    frame_size_bytes = int((write_rate_mbps * 1024 * 1024) / 8 / 30)  # 每帧字节数
    write_interval = 1 / 30  # 30fps，每帧间隔
    total_frames = int(duration_seconds * 30)  # 总帧数
    
    # 创建测试文件
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    test_file = os.path.join(test_dir, f"test_{usb_version}_{timestamp}.dat")
    
    # 准备测试数据块（模拟视频帧）
    data_block = generate_test_data(frame_size_bytes)
    
    # 性能记录
    frame_times = []  # 每帧写入时间
    write_speeds = []  # 每秒写入速度
    last_second_bytes = 0  # 每秒写入字节数
    error_count = 0  # 错误次数
    
    print(f"\n=== 开始 {usb_version} 接口测试 ===")
    print(f"测试文件: {test_file}")
    print(f"目标写入速率: {write_rate_mbps}Mbps ({frame_size_bytes/1024:.2f}KB/帧 @ 30fps)")
    print(f"测试时长: {duration_seconds}秒")
    
    # 开始测试
    start_time = time.perf_counter()
    last_frame_time = start_time
    last_second_time = start_time
    next_frame_time = start_time + write_interval
    
    with open(test_file, "wb") as f:
        for frame_num in range(total_frames):
            # 等待下一帧时间点（模拟实时录制）
            current_time = time.perf_counter()
            wait_time = next_frame_time - current_time
            if wait_time > 0:
                time.sleep(wait_time)
            
            # 记录实际写入时间
            frame_start_time = time.perf_counter()
            
            # 写入数据（模拟视频帧）
            try:
                f.write(data_block)
                f.flush()  # 强制刷新到磁盘
            except Exception as e:
                print(f"错误: 写入失败 (帧#{frame_num}): {str(e)}")
                error_count += 1
            
            # 记录帧写入时间
            frame_end_time = time.perf_counter()
            frame_times.append(frame_end_time - frame_start_time)
            
            # 计算每秒写入速度
            last_second_bytes += frame_size_bytes
            if frame_end_time - last_second_time >= 1.0:
                speed_mbps = (last_second_bytes * 8) / (1024 * 1024)
                write_speeds.append(speed_mbps)
                last_second_bytes = 0
                last_second_time = frame_end_time
            
            # 更新下一帧时间
            next_frame_time += write_interval
            
            # 显示进度
            if frame_num % 30 == 0:  # 每秒更新一次进度
                elapsed = frame_end_time - start_time
                progress = (frame_num / total_frames) * 100
                avg_frame_time = sum(frame_times[-30:]) / min(30, len(frame_times))
                print(f"\r进度: {progress:.1f}% | 已写入: {frame_num//30}秒 | 平均帧写入时间: {avg_frame_time*1000:.2f}ms", end="")
    
    # 测试完成
    end_time = time.perf_counter()
    total_duration = end_time - start_time
    
    # 计算统计数据
    avg_frame_time = sum(frame_times) / len(frame_times)
    max_frame_time = max(frame_times)
    min_frame_time = min(frame_times)
    avg_speed = sum(write_speeds) / len(write_speeds) if write_speeds else 0
    max_speed = max(write_speeds) if write_speeds else 0
    min_speed = min(write_speeds) if write_speeds else 0
    
    # 输出结果
    print(f"\n=== {usb_version} 测试结果 ===")
    print(f"总耗时: {total_duration:.2f}秒")
    print(f"写入总数据量: {frame_size_bytes * total_frames / (1024*1024):.2f}MB")
    print(f"帧写入时间: 平均 {avg_frame_time*1000:.2f}ms, 最小 {min_frame_time*1000:.2f}ms, 最大 {max_frame_time*1000:.2f}ms")
    print(f"写入速度: 平均 {avg_speed:.2f}Mbps, 最小 {min_speed:.2f}Mbps, 最大 {max_speed:.2f}Mbps")
    print(f"写入错误次数: {error_count}")
    
    # 评估是否满足录制要求
    if max_frame_time > write_interval * 1.5:
        print("警告: 存在写入延迟峰值，可能导致视频丢帧或卡顿")
    if error_count > 0:
        print("警告: 写入过程中出现错误，视频文件可能损坏")
    
    # 保存测试报告
    report_file = os.path.join(test_dir, f"test_{usb_version}_{timestamp}_report.txt")
    with open(report_file, "w") as f:
        f.write(f"测试时间: {datetime.datetime.now()}\n")
        f.write(f"USB接口: {usb_version}\n")
        f.write(f"目标写入速率: {write_rate_mbps}Mbps\n")
        f.write(f"总耗时: {total_duration:.2f}秒\n")
        f.write(f"写入总数据量: {frame_size_bytes * total_frames / (1024*1024):.2f}MB\n")
        f.write(f"帧写入时间: 平均 {avg_frame_time*1000:.2f}ms, 最小 {min_frame_time*1000:.2f}ms, 最大 {max_frame_time*1000:.2f}ms\n")
        f.write(f"写入速度: 平均 {avg_speed:.2f}Mbps, 最小 {min_speed:.2f}Mbps, 最大 {max_speed:.2f}Mbps\n")
        f.write(f"写入错误次数: {error_count}\n")
    
    return {
        "usb_version": usb_version,
        "avg_speed": avg_speed,
        "max_speed": max_speed,
        "min_speed": min_speed,
        "avg_frame_time": avg_frame_time,
        "max_frame_time": max_frame_time,
        "error_count": error_count
    }


def main():
    parser = argparse.ArgumentParser(description="U盘视频录制性能测试工具")
    parser.add_argument("--usb3-dir", required=True, help="USB3.0接口测试目录（U盘根目录）")
    parser.add_argument("--usb2-dir", required=True, help="USB2.0接口测试目录（U盘根目录）")
    parser.add_argument("--duration", type=int, default=300, help="测试持续时间（秒），默认300秒")
    parser.add_argument("--rate", type=int, default=20, help="目标写入速率（Mbps），默认20Mbps")
    args = parser.parse_args()
    
    # 验证目录
    if not os.path.isdir(args.usb3_dir):
        print(f"错误: USB3.0测试目录不存在 - {args.usb3_dir}")
        return
    
    if not os.path.isdir(args.usb2_dir):
        print(f"错误: USB2.0测试目录不存在 - {args.usb2_dir}")
        return
    
    # 执行测试
    usb3_result = test_recording("USB3.0", args.usb3_dir, args.duration, args.rate)
    print("\n\n请将U盘从USB3.0接口拔出，插入USB2.0接口，然后按Enter继续...")
    input()
    usb2_result = test_recording("USB2.0", args.usb2_dir, args.duration, args.rate)
    
    # 对比结果
    print("\n\n=== 测试对比结果 ===")
    print(f"目标写入速率: {args.rate}Mbps")
    
    print("\n写入速度对比:")
    print(f"USB3.0: 平均 {usb3_result['avg_speed']:.2f}Mbps, 最大 {usb3_result['max_speed']:.2f}Mbps")
    print(f"USB2.0: 平均 {usb2_result['avg_speed']:.2f}Mbps, 最大 {usb2_result['max_speed']:.2f}Mbps")
    print(f"USB3.0速度优势: {usb3_result['avg_speed']/usb2_result['avg_speed']:.2f}倍")
    
    print("\n帧写入时间对比:")
    print(f"USB3.0: 平均 {usb3_result['avg_frame_time']*1000:.2f}ms, 最大 {usb3_result['max_frame_time']*1000:.2f}ms")
    print(f"USB2.0: 平均 {usb2_result['avg_frame_time']*1000:.2f}ms, 最大 {usb2_result['max_frame_time']*1000:.2f}ms")
    
    print("\n错误情况:")
    print(f"USB3.0: 错误次数 {usb3_result['error_count']}")
    print(f"USB2.0: 错误次数 {usb2_result['error_count']}")
    
    # 评估是否适合视频录制
    print("\n=== 适用性评估 ===")
    if usb3_result['avg_speed'] >= args.rate * 0.9 and usb3_result['error_count'] == 0:
        print("USB3.0: ✅ 适合目标码率的视频录制")
    else:
        print("USB3.0: ❌ 不适合目标码率的视频录制")
    
    if usb2_result['avg_speed'] >= args.rate * 0.9 and usb2_result['error_count'] == 0:
        print("USB2.0: ✅ 适合目标码率的视频录制")
    else:
        print("USB2.0: ❌ 不适合目标码率的视频录制")


if __name__ == "__main__":
    main()